// SPDX-License-Identifier: GPL-2.0
#define KMSG_COMPONENT "IPVS"
#define pr_fmt(fmt) KMSG_COMPONENT ": " fmt

#include <linux/ip.h>
#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/skbuff.h>
#include <linux/rbtree.h>
#include <linux/crypto.h>
#include <linux/scatterlist.h>
#include "./ip_vs_fnat.h"


#define IP_VS_SH_TAB_SIZE               (1 << 8)
#define IP_VS_DEST_USED_IN_RBTREE	0x8000	/* dest is used in rbtree*/

#define IP_VS_CONSH_LEN			70


enum {
	IP_VS_FLAG_UNUSED = 0,
	IP_VS_FLAG_ADD,
	IP_VS_FLAG_INUSED
};

struct ip_vs_consh_tree {
	struct rb_root rbroot;
	struct list_head node_list;
	unsigned int num_nodes;
	u32 (*hashfunc_handler)(const char *);
};

/* nodes structure */
struct ip_vs_consh_node {
	struct list_head n_list;
	u16 af;			/* address family */
	union nf_inet_addr addr;	/* IP address of the server */
	__be16 port;		/* port number of the server */
	/* node name or some thing identifies the node */
	char iden[IP_VS_CONSH_LEN];
	u_int replicas; /* number of replica virtual nodes */
	u_int flag;
	__u32 vni;
	void *data;
};

struct ip_vs_consh_entry {
	struct rb_node rbnode;
	struct ip_vs_consh_node *node;
	long key;
};

u32 ip_vs_consh_jhash(const char *instr)
{
	return jhash(instr, strlen(instr), 0);
}

static inline void ip_vs_consh_iden(struct ip_vs_consh_node *node, int idx)
{
	memset(node->iden, 0, IP_VS_CONSH_LEN);
	if (node->af == AF_INET) {
		snprintf(node->iden, IP_VS_CONSH_LEN,
		"%u-%pI4-%u-%03d", node->vni, &node->addr.ip,
		ntohs(node->port), idx);
	} else {
		snprintf(node->iden, IP_VS_CONSH_LEN,
		"%u-%pI6-%u-%03d", node->vni, &node->addr.ip6,
		ntohs(node->port), idx);
	}
}

static struct ip_vs_consh_entry *ipvs_rb_lookup(struct rb_root *root, long key)
{
	struct rb_node *tmp = root->rb_node;
	struct ip_vs_consh_entry *entry = NULL;
	struct rb_node *node = NULL;

	while (tmp != NULL) {
		entry = rb_entry(tmp, struct ip_vs_consh_entry, rbnode);
		if (entry->key > key) {
			node = tmp; /* set this can get better balance */
			tmp = tmp->rb_left;
		} else if (entry->key < key) {
			tmp = tmp->rb_right;
		} else
			return entry;
	}

	if (node != NULL)
		entry = rb_entry(node, struct ip_vs_consh_entry, rbnode);
	return ((entry != NULL) ? entry :  rb_entry(rb_first(root), struct ip_vs_consh_entry, rbnode));
}

static struct ip_vs_consh_entry *ipvs_rb_search(struct rb_root *root, long key)
{
	struct rb_node *node = root->rb_node;
	struct ip_vs_consh_entry *entry;

	while (node != NULL) {
		entry = rb_entry(node, struct ip_vs_consh_entry, rbnode);
		if (entry->key > key)
			node = node->rb_left;
		else if (entry->key < key)
			node = node->rb_right;
		else
			return entry;
	}
	return NULL;
}

static int ipvs_rb_insert(struct rb_root *root, struct ip_vs_consh_entry *entry,
			struct ip_vs_consh_entry **dupentry)
{
	struct rb_node **link = &root->rb_node, *parent = NULL;
	struct ip_vs_consh_entry *myentry;

	while (*link) {
		parent = *link;
		myentry = rb_entry(parent, struct ip_vs_consh_entry, rbnode);
		if (myentry->key > entry->key)
			link = &(*link)->rb_left;
		else if (myentry->key < entry->key)
			link = &(*link)->rb_right;
		else {
			*dupentry = myentry;
			return -EEXIST;
		}
	}
	rb_link_node(&entry->rbnode, parent, link);
	rb_insert_color(&entry->rbnode, root);
	return 0;
}

static void ip_vs_del_node(struct ip_vs_consh_tree *tree, struct ip_vs_consh_node *node,
			    struct ip_vs_service *svc)
{
	u32 hash;
	int i;
	struct ip_vs_consh_entry *entry;

	for (i = 0; i < node->replicas; i++) {
		ip_vs_consh_iden(node, i);
		hash = tree->hashfunc_handler(node->iden);
		entry = ipvs_rb_search(&tree->rbroot, hash);
		if (entry != NULL) {
			rb_erase(&entry->rbnode, &tree->rbroot);
			kfree(entry);
			entry = NULL;
		} else
			pr_err("%s not searched\n", node->iden);
	}
	((struct ip_vs_dest *)(node->data))->flags &= ~IP_VS_DEST_USED_IN_RBTREE;
	refcount_dec(&((struct ip_vs_dest *)(node->data))->refcnt);
	node->data = NULL;
	list_del(&node->n_list);
	tree->num_nodes--;
	node->flag = IP_VS_FLAG_UNUSED;
	IP_VS_DBG_BUF(6, "del  %s %u\n",
		IP_VS_DBG_ADDR(node->af, &node->addr),
			       ntohs(node->port));
	kfree(node);
	node = NULL;
}

static int
ip_vs_consh_new_node(struct ip_vs_consh_tree *tree, struct ip_vs_dest *dest)
{
	u32 hash;
	int i, ret;
	struct ip_vs_consh_node *new_node;
	struct ip_vs_consh_entry *entry, *dupentry = NULL;

	if (dest->flags & IP_VS_DEST_USED_IN_RBTREE)
		return 0;

	dest->flags |= IP_VS_DEST_USED_IN_RBTREE;
	new_node = kzalloc(sizeof(struct ip_vs_consh_node), GFP_ATOMIC);
	if (!new_node) {
		pr_err(" kzalloc entry failed in consh\n");
		return -1;
	}
	ip_vs_addr_copy(dest->af, &new_node->addr, &dest->addr);
	new_node->af = dest->af;
	new_node->port = dest->port;
	new_node->vni = dest->vni.out_vni;
	new_node->replicas = IP_VS_SH_TAB_SIZE;
	new_node->flag = IP_VS_FLAG_ADD;
	new_node->data = (void *)dest;

	/* using in consh tree, inc dest refcnt */
	refcount_inc(&((struct ip_vs_dest *)(new_node->data))->refcnt);
	list_add(&new_node->n_list, &tree->node_list);
	tree->num_nodes++;
	for (i = 0; i < new_node->replicas; i++) {
		ip_vs_consh_iden(new_node, i);
		hash = tree->hashfunc_handler(new_node->iden);
		entry = kzalloc(sizeof(struct ip_vs_consh_entry), GFP_ATOMIC);
		if (!entry) {
			pr_err(" kzalloc entry failed in consh\n");
			new_node->replicas = i;
			ip_vs_del_node(tree, new_node, dest->svc);
			return -1;
		}
		entry->key = hash;
		entry->node = new_node;
		ret = ipvs_rb_insert(&tree->rbroot, entry, &dupentry);
		if (ret) {
			entry->node = NULL;
			kfree(entry);
			entry = NULL;
		}
	}
	IP_VS_DBG_BUF(6, "add %s vni:%d,port:%d\n",
		IP_VS_DBG_ADDR(new_node->af, &new_node->addr),
		new_node->vni, ntohs(new_node->port));

	return 0;
}

static void
ip_vs_consh_init_step(struct ip_vs_consh_tree *tree, struct ip_vs_service *svc)
{
	struct ip_vs_dest *dest;

	if (!svc->num_dests)
		return;

	list_for_each_entry_rcu(dest, &svc->destinations, n_list) {
		ip_vs_consh_new_node(tree, dest);
	}
}

static int ip_vs_consh_assign(struct ip_vs_consh_tree *tree,
		struct ip_vs_service *svc, struct ip_vs_dest *dest)
{
	struct ip_vs_consh_node *node = NULL;
	struct ip_vs_consh_node *match_node = NULL;

	if (!tree)
		return 0;

	if (!dest) {
		/* init step */
		ip_vs_consh_init_step(tree, svc);
		return 0;
	}

	if (svc->num_dests >= tree->num_nodes) {/* add node */
		/* Edit dest case, do nothing. */
		if (dest->flags & IP_VS_DEST_USED_IN_RBTREE)
			return 0;

		ip_vs_consh_new_node(tree, dest);
	} else {/* del */
		list_for_each_entry(node, &tree->node_list, n_list) {
			if ((node->port == dest->port) &&
			    ip_vs_addr_equal(node->af, &node->addr, &dest->addr) &&
			    (node->af == dest->af) &&
			    (node->vni == dest->vni.out_vni)) {
				match_node = node;
				break;
			}
		}

		if (match_node && match_node->flag == IP_VS_FLAG_ADD)
			ip_vs_del_node(tree, match_node, svc);
		else
			pr_err("Scheduler consh must be in a mess.\n");
	}

	return 0;
}

static int ip_vs_consh_init_svc(struct ip_vs_service *svc)
{
	struct ip_vs_consh_tree *tree;

	/* allocate the SH table for this service */
	tree = kzalloc(sizeof(struct ip_vs_consh_tree), GFP_ATOMIC);
	if (tree == NULL) {
		pr_err("%s(): no memory\n", __func__);
		return -ENOMEM;
	}

	tree->rbroot = RB_ROOT;
	tree->num_nodes = 0;
	tree->hashfunc_handler = ip_vs_consh_jhash;
	INIT_LIST_HEAD(&tree->node_list);

	svc->sched_data = tree;

	IP_VS_DBG(6, "conSH hash table (memory=%zd bytes) allocated for "
		  "current service\n", sizeof(struct ip_vs_consh_tree));

	/* assign the hash buckets with the updated service */
	ip_vs_consh_assign(tree, svc, NULL);

	return 0;
}
static void ip_vs_consh_done_svc(struct ip_vs_service *svc)
{
	int i;
	struct ip_vs_consh_tree *tree = svc->sched_data;
	struct ip_vs_consh_node *node, *node_next;
	struct ip_vs_consh_entry *entry;
	u32 hash;

	list_for_each_entry_safe(node, node_next, &tree->node_list, n_list) {
		for (i = 0; i < node->replicas; i++) {
			ip_vs_consh_iden(node, i);
			hash = tree->hashfunc_handler(node->iden);
			entry = ipvs_rb_search(&tree->rbroot, hash);
			if (entry != NULL) {
				rb_erase(&entry->rbnode, &tree->rbroot);
				kfree(entry);
				entry = NULL;
			}
		}
		list_del(&node->n_list);
		if (tree->num_nodes == 0)
			pr_err("%s(): tree node count error \n", __func__);
		tree->num_nodes--;
		((struct ip_vs_dest *)(node->data))->flags &= ~IP_VS_DEST_USED_IN_RBTREE;
		refcount_dec(&((struct ip_vs_dest *)(node->data))->refcnt);
		node->data = NULL;
		node->flag = IP_VS_FLAG_UNUSED;
		IP_VS_DBG_BUF(6, "del all %s vni:%d,port:%u\n",
			IP_VS_DBG_ADDR(node->af, &node->addr),
				node->vni, ntohs(node->port));
		kfree(node);
		node = NULL;
	}

	spin_lock_bh(&svc->sched_lock);
	/* release the table itself */
	kfree(svc->sched_data);
	svc->sched_data = NULL;
	spin_unlock_bh(&svc->sched_lock);
}

static int ip_vs_consh_update_svc(struct ip_vs_service *svc, struct ip_vs_dest *dest)
{
	spin_lock_bh(&svc->sched_lock);
	ip_vs_consh_assign((struct ip_vs_consh_tree *)svc->sched_data, svc, dest);
	spin_unlock_bh(&svc->sched_lock);

	return 0;
}

/*
 *      If the dest flags is set with IP_VS_DEST_F_OVERLOAD,
 *      consider that the server is overloaded here.
 */
static inline int is_overloaded(struct ip_vs_dest *dest)
{
	return dest->flags & IP_VS_DEST_F_OVERLOAD;
}

/*
 *      Get ip_vs_dest associated with supplied parameters.
 */
static inline struct ip_vs_dest *ip_vs_consh_get(int af,
						 struct ip_vs_service *svc,
						 const union nf_inet_addr *addr)
{
	char buf[IP_VS_CONSH_LEN] = {0};
	u32 hash;
	struct ip_vs_consh_entry *entry;
	struct ip_vs_consh_tree *tree = NULL;
	struct ip_vs_dest *dest = NULL;

	if (af == AF_INET6) {
		snprintf(buf, sizeof(buf), "%pI6", &addr->ip6);
	} else {
		snprintf(buf, sizeof(buf), "%pI4", &addr->ip);
	}
	IP_VS_DBG(6, "get cons ip address %s\n", buf);

	spin_lock_bh(&svc->sched_lock);
	tree = (struct ip_vs_consh_tree *)svc->sched_data;
	if (!tree)
		goto tree_err;

	hash = tree->hashfunc_handler(buf);
	entry = ipvs_rb_lookup(&tree->rbroot, hash);
	if (entry && entry->node)
		dest = (struct ip_vs_dest *)entry->node->data;

tree_err:
	spin_unlock_bh(&svc->sched_lock);
	return dest;
}
/*
 *      Source Hashing scheduling
 */
static struct ip_vs_dest *ip_vs_consh_schedule(struct ip_vs_service *svc,
					    const struct sk_buff *skb, struct ip_vs_iphdr *iph)
{
	struct ip_vs_dest *dest;

	IP_VS_DBG(6, "ip_vs_sh_schedule(): Scheduling...\n");
	dest = ip_vs_consh_get(svc->af, svc, &iph->saddr);
	if (!dest || !(dest->flags & IP_VS_DEST_F_AVAILABLE)
	    || atomic_read(&dest->weight) <= 0 || is_overloaded(dest)) {
		IP_VS_ERR_RL("CONSH: no destination available\n");
		return NULL;
	}
	IP_VS_DBG_BUF(6, "CONSH: source IP address %s --> server %s:%u\n",
			IP_VS_DBG_ADDR(svc->af, &iph->saddr),
			IP_VS_DBG_ADDR(svc->af, &dest->addr),
			ntohs(dest->port));

	return dest;
}

/*
 *      IPVS SH Scheduler structure
 */
static struct ip_vs_scheduler ip_vs_consh_scheduler = {
	.name = "consh",
	.refcnt = ATOMIC_INIT(0),
	.module = THIS_MODULE,
	.n_list = LIST_HEAD_INIT(ip_vs_consh_scheduler.n_list),
	.init_service = ip_vs_consh_init_svc,
	.done_service = ip_vs_consh_done_svc,
	.add_dest = ip_vs_consh_update_svc,
	.del_dest = ip_vs_consh_update_svc,
	.upd_dest = ip_vs_consh_update_svc,
	.schedule = ip_vs_consh_schedule,
};

static int __init ip_vs_consh_init(void)
{
	return register_ip_vs_scheduler(&ip_vs_consh_scheduler);
}

static void __exit ip_vs_consh_cleanup(void)
{
	unregister_ip_vs_scheduler(&ip_vs_consh_scheduler);
}

module_init(ip_vs_consh_init);
module_exit(ip_vs_consh_cleanup);
MODULE_LICENSE("GPL v2");

